You write custom CUDA kernels to replace pytorch operators in given architecture to get speedups. You have complete freedom to choose set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.

**SPECIAL INSTRUCTIONS FOR COSINE DISTANCE + SOFTMAX FUSION:**

When implementing Cosine Distance + Softmax fusion, you MUST implement the following optimized strategy:

1. **THREE-PHASE COMPUTATION**: Implement as three separate kernels for maximum precision:
   - Phase 1: Compute vector norms using warp-level reduction
   - Phase 2: Compute cosine similarities using pre-computed norms
   - Phase 3: Apply softmax weighting to cosine distances

2. **PURE CUDA IMPLEMENTATION**: Use only pure CUDA functions, no PyTorch internal functions:
   - Use only CUDA built-in functions: sqrtf, expf, fmaxf, __shfl_down_sync
   - Use only CUDA memory management: extern __shared__, reinterpret_cast
   - No PyTorch tensor operations inside CUDA kernels

3. **WARP-LEVEL OPTIMIZATION**: Use warp-level processing for maximum performance:
   - Each block processes one sample from the batch
   - Use 8 warps per block (256 threads) for optimal GPU utilization
   - Use __shfl_down_sync for efficient warp-level reduction

4. **MEMORY COALESCING WITH float4**: Use vectorized memory access:
   - Load 4 elements at once using float4 for memory efficiency
   - Process 4 elements per thread to maximize throughput
   - Ensure proper alignment for coalesced memory access

5. **THREE-PHASE SOFTMAX ALGORITHM**: Implement numerically stable softmax:
cpp
// Phase 1: Find maximum value across all samples
float max_val = -FLT_MAX;
for (int idx = tid; idx < batch_size; idx += blockDim.x) {
    max_val = fmaxf(max_val, cosine_sims[idx] / temperature);
}

// Phase 2: Compute exp values and sum
float sum_exp = 0.0f;
for (int idx = tid; idx < batch_size; idx += blockDim.x) {
    float exp_val = expf(cosine_sims[idx] / temperature - global_max);
    sum_exp += exp_val;
    shared_exp[idx] = exp_val;
}

// Phase 3: Compute softmax weights and weighted distances
for (int idx = tid; idx < batch_size; idx += blockDim.x) {
    float softmax_weight = shared_exp[idx] / global_sum;
    float cosine_dist = 1.0f - cosine_sims[idx];
    weighted_distances[idx] = softmax_weight * cosine_dist;
}



6. **SHARED MEMORY PATTERN**: Use efficient shared memory organization:
cpp
// For norm computation
extern __shared__ float shared_norms[];
if (lane_id == 0) {
    shared_norms[warp_id] = warp_norm_sq;
}

// For cosine similarity
extern __shared__ float shared_cosine[];
if (lane_id == 0) {
    shared_cosine[warp_id] = warp_dot_product;
}

// For softmax computation
extern __shared__ float shared_softmax[];
shared_softmax[32 + idx] = exp_val; // Store exp values
shared_softmax[warp_id] = max_val;   // Store reduction results



7. **BLOCK CONFIGURATION**: Use optimal settings:
   - Block size: 256 threads (8 warps)
   - Shared memory: Dynamic allocation based on batch size
   - One block per sample for norm and cosine computation
   - Single block for softmax computation (global operation)

8. **PRECISION REQUIREMENTS**: Ensure exact mathematical alignment:
   - Cosine similarity: cosine_sim = (x·y) / (||x||·||y||)
   - Softmax: softmax_weight = exp(cosine_sim/T) / sum(exp(cosine_sim/T))
   - Weighted distance: weighted_dist = softmax_weight * (1 - cosine_sim)
   - Verify with torch.allclose(rtol=1e-03, atol=1e-6)

9. **FUNCTION SIGNATURE**: The main CUDA function must accept all parameters:
cpp
torch::Tensor cosinedistance_softmax_cuda(
    torch::Tensor x,
    torch::Tensor y,
    float eps,
    float temperature
)



10. **MATHEMATICAL FORMULAS**: Implement exact mathematical operations:
    - Cosine Similarity: cos_sim = (x·y) / (||x||·||y||)
    - Softmax: softmax_i = exp(cos_sim_i/T) / Σ_j exp(cos_sim_j/T)
    - Cosine Distance: cos_dist_i = 1 - cos_sim_i
    - Weighted Distance: weighted_dist_i = softmax_i × cos_dist_i

11. **PYTHON CALLING CONVENTION**: The ModelNew forward method must pass parameters correctly:
python
def forward(self, x, y):
    return self.cosinedistance_softmax.cosinedistance_softmax_cuda(x, y, self.eps, self.temperature)



Here's the target architecture to optimize:

python
import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
"""
CosineDistance + Softmax fusion implementation.
Computes cosine similarities, applies softmax weighting, then calculates weighted cosine distances.
"""
def __init__(self, eps=1e-8, temperature=1.0):
    super(Model, self).__init__()
    self.eps = eps
    self.temperature = temperature

def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """
    Compute CosineDistance + Softmax weighted fusion.

    Args:
        x (torch.Tensor): First set of vectors [batch_size, feature_dim]
        y (torch.Tensor): Second set of vectors [batch_size, feature_dim]

    Returns:
        torch.Tensor: Weighted cosine distances [batch_size]
    """
    # Compute cosine similarities
    cosine_sim = F.cosine_similarity(x, y, dim=1, eps=self.eps)
    
    # Apply softmax weighting with temperature
    softmax_weights = F.softmax(cosine_sim / self.temperature, dim=0)
    
    # Compute weighted cosine distances
    cosine_dist = 1.0 - cosine_sim
    weighted_distances = softmax_weights * cosine_dist
    
    return weighted_distances

batch_size = 256
feature_dim = 512

def get_inputs():
    # Generate two sets of vectors
    x = torch.randn(batch_size, feature_dim)
    y = torch.randn(batch_size, feature_dim)
    return [x, y]

def get_init_inputs():
    return [1e-8, 1.0]  # eps and temperature values



**EXPECTED OUTPUT STRUCTURE**:
Generate two files:
1. `cosinedistance_softmax_cudacode.py` - Contains ModelNew class with CosineDistance+Softmax fusion using pure CUDA
2. `cosinedistance_softmax_torchcode.py` - Contains the reference PyTorch implementation

**KEY REQUIREMENTS**:
- The CUDA implementation must use pure CUDA functions only
- Must implement three-phase computation: norms → cosine similarities → softmax weighting
- Must use warp-level optimization for maximum performance
- Must use float4 vectorized memory access for coalescing
- Must use three-phase softmax algorithm for numerical stability
- Must handle arbitrary tensor shapes (not just fixed dimensions)
- Must maintain mathematical precision with PyTorch implementation
- Must use optimal block configuration (256 threads, 8 warps)
- Expected speedup: 2.0-3.0x over PyTorch baseline
- Must use fast math optimizations for better performance
- Must be robust and handle edge cases properly
- Must use only pure CUDA functions (no PyTorch internal functions)
- Must use sqrtf, expf, fmaxf for mathematical operations
- Must implement exact mathematical formulas for cosine similarity and softmax
- Must pass eps and temperature parameters correctly from Python to CUDA
- Must use Python float syntax (0.01) not C++ syntax (0.01f) in Python code
- Must implement three separate kernels for maximum precision
- Must use shared memory efficiently for warp-level reductions
- Must ensure global softmax computation across all samples
